#################################
##' Code for:
##' Title: "Effective climate clubs require ambition, leverage, and insulation"
##' Author: Sam S. Rowan
##' Contact: sam.rowan [at] concordia.ca
##' Date: 2023-10-19
#################################
##' This script:
##' Distribution of membership varying theta 
#################################

#### Initialize environment ####

## Packages
library(tidyverse)
library(truncnorm)
library(knitr)

## Working directory
# setwd("/Users/srowa/Dropbox/Projects/climate_clubs_trade")
# -- Set working directory to replication folder

## Define the logistic function
logistic_function <- function(x, k, a) {
  return(1 / (1 + exp(-k * (x - a))))
}

## Plot settings for later
theme_clubs <-  theme(legend.position = "bottom", 
                      panel.background = element_rect(fill=NA),
                      panel.grid.major.x = element_blank(),
                      panel.grid.major.y = element_blank(),
                      panel.grid.minor.y =  element_blank(),
                      panel.grid.minor.x =  element_blank(),
                      panel.ontop = FALSE,
                      axis.line = element_line(linewidth = 1, colour = "black"),
                      axis.text = element_text(size = 16),
                      legend.text = element_text(size = 12),
                      legend.title = element_text(size = 16),
                      axis.title = element_text(size = 16))


#### Define climate model parameters ####

## Define number of states
num_states <- 150

## Draw their climate ideal points
set.seed(202109)
climate_ideal <- cbind.data.frame(
  climate_ideal = runif(n = num_states, min = 0, max = 1)) %>% 
  arrange(climate_ideal) %>% 
  mutate(id = seq(from = 1, to = num_states, by = 1))
# climate_ideal

####' ---- NOTE: Code block below runs for a couple hours on 2021 iMac
####' ---- The simulation output is saved as 'output/sims_broader_deeper_IxI.csv'
####' ---- Skip running the simulation and load this in directly at around line 300
####' ---- And build figures from there

#### Simulate trade ties 500 times, analyze across theta ####

## Create loop counters
set.seed(202109)
seeds_vector <- as.integer(runif(n = 100, min = 1, max = 1e6))
theta_vector <- seq(0.05, 0.95, 0.01)

## Create objects to hold outputs
# Count membership for each value of theta and each draw of trade ties
length_members_independent <- matrix(NA, nrow = length(theta_vector), ncol = length(seeds_vector))
length_members_positive <- matrix(NA, nrow = length(theta_vector), ncol = length(seeds_vector))
length_members_negative <- matrix(NA, nrow = length(theta_vector), ncol = length(seeds_vector))
# Count emissions for each value of theta and each draw of trade ties
length_emissions_independent <- matrix(NA, nrow = length(theta_vector), ncol = length(seeds_vector))
length_emissions_positive <- matrix(NA, nrow = length(theta_vector), ncol = length(seeds_vector))
length_emissions_negative <- matrix(NA, nrow = length(theta_vector), ncol = length(seeds_vector))

## Simulation model
start_time <- Sys.time()
# Declare values of theta
for (theta_draw in theta_vector) {
  for (seed in seeds_vector) {
    # Declare seed
    set.seed(seed)
    theta_df <-  cbind.data.frame(
      # Add IDs
      id_a = rep(climate_ideal$id, times = num_states),
      id_b = rep(climate_ideal$id, each = num_states),
      # Bind climate ideal points into dyadic setup
      climate_ideal_a = rep(climate_ideal$climate_ideal, times = num_states),
      climate_ideal_b = rep(climate_ideal$climate_ideal, each = num_states)) %>% 
      # Measure dyadic ideal point distance
      mutate(ideal_point_distance = abs(climate_ideal_a - climate_ideal_b)) %>% 
      # --** Random component **--
      # Draw trade data based on ideal point distance
      mutate(trade_independent = rtruncnorm(n = num_states^2,
                                            # a = 0, b = 1,
                                            mean = 0,
                                            sd = 0.5),
             trade_positive = rtruncnorm(n = num_states^2,
                                         # a = 0, b = 1,
                                         mean = (1-ideal_point_distance), # negative of difference, implies trade more when ideal points closer
                                         sd = 0.5),
             trade_negative = rtruncnorm(n = num_states^2,
                                         # a = 0, b = 1,
                                         mean = ideal_point_distance, # positive of difference, implies trade less when ideal points closer
                                         sd = 0.5)) %>% 
      # Use logistic function to bound trade ties
      mutate_at(vars(trade_independent:trade_negative),
                ~logistic_function(x = ., k = 10, a = 0.75)) %>% 
      # Remove data for faux dyads -- where id_a == id_b
      mutate_at(vars(ideal_point_distance,
                     trade_independent, trade_positive, trade_negative),
                ~replace(., id_a == id_b, NA)) %>% 
      # Get each state's total trade
      group_by(id_a) %>% 
      mutate(sum_trade_independent = sum(trade_independent, na.rm = T),
             sum_trade_positive = sum(trade_positive, na.rm = T),
             sum_trade_negative = sum(trade_negative, na.rm = T)) %>% 
      # Standardize
      rowwise() %>% 
      mutate(share_dyadic_independent = trade_independent / sum_trade_independent,
             share_dyadic_positive = trade_positive / sum_trade_positive,
             share_dyadic_negative = trade_negative / sum_trade_negative) %>% 
      select(id_a, id_b, climate_ideal_a, climate_ideal_b, ideal_point_distance,
             share_dyadic_independent, share_dyadic_positive, share_dyadic_negative) %>% 
      # Create a club's common climate goal
      mutate(theta = theta_draw,
             core_member_a = ifelse(climate_ideal_a >= theta, 1, 0),
             core_member_b = ifelse(climate_ideal_b >= theta, 1, 0)) %>% 
      # Flag trade to club members
      mutate(trade_to_club_independent = ifelse(core_member_b == 1, 
                                                share_dyadic_independent, NA),
             trade_to_club_positive = ifelse(core_member_b == 1, 
                                             share_dyadic_positive, NA),
             trade_to_club_negative = ifelse(core_member_b == 1, 
                                             share_dyadic_negative, NA)) %>% 
      # Get share of trade to all club members
      group_by(id_a) %>% 
      mutate(share_trade_club_independent = sum(trade_to_club_independent, na.rm = T),
             share_trade_club_positive = sum(trade_to_club_positive, na.rm = T),
             share_trade_club_negative = sum(trade_to_club_negative, na.rm = T)) %>% 
      ungroup() %>% 
      select(-starts_with("trade_to_club")) %>% 
      # Collapse the data to just _a's perspective, to calculate utilities
      select(id_a, climate_ideal_a, theta, core_member_a,
             starts_with("share_trade_club")) %>% 
      distinct() %>% # rows = 150
      # Calculate utilities to membership
      mutate(beta = 1) %>% # Define weights for climate versus trade
      rowwise() %>% 
      # Get most constrained core club member
      mutate(min_trade_exposure_independent = ifelse(
        core_member_a == 1, min(share_trade_club_independent), NA),
        min_trade_exposure_positive = ifelse(
          core_member_a == 1, min(share_trade_club_positive ), NA),
        min_trade_exposure_negative = ifelse(
          core_member_a == 1, min(share_trade_club_negative), NA)) %>% 
      ungroup() %>% 
      # Define that constraint as lambda
      mutate(lambda_independent = min(min_trade_exposure_independent, na.rm = T),
             lambda_positive = min(min_trade_exposure_positive, na.rm = T),
             lambda_negative = min(min_trade_exposure_negative, na.rm = T)) %>% 
      # Calculate utilities for the independent world
      # Utility to join
      mutate(utility_join_independent = (-abs(climate_ideal_a - theta)) + # Climate
               (1 - share_trade_club_independent) * (1 - lambda_independent)*beta + # Non-club trade
               (share_trade_club_independent*beta) - # Club trade
               (theta - lambda_independent)) %>%  # Aversion
      # Utility to abstain
      mutate(utility_abstain_independent = (-climate_ideal_a) + # Climate
               (1-share_trade_club_independent)*beta + # Non-club trade
               (share_trade_club_independent*(1-lambda_independent)*beta) + # Club trade
               (0)) %>%  # Aversion
      # Net, membership
      mutate(net_utility_independent = utility_join_independent - utility_abstain_independent,
             membership_independent = ifelse(net_utility_independent >= 0, 1, 0)) %>% 
      # Calculate utilities for the positive world
      # Utility to join
      mutate(utility_join_positive = (-abs(climate_ideal_a - theta)) + # Climate
               (1 - share_trade_club_positive) * (1 - lambda_positive)*beta + # Non-club trade
               (share_trade_club_positive*beta) - # Club trade
               (theta - lambda_positive)) %>%  # Aversion
      # Utility to abstain
      mutate(utility_abstain_positive = (-climate_ideal_a) + # Climate
               (1-share_trade_club_positive)*beta + # Non-club trade
               (share_trade_club_positive*(1-lambda_positive)*beta) + # Club trade
               (0)) %>%  # Aversion
      # Net, membership
      mutate(net_utility_positive = utility_join_positive - utility_abstain_positive,
             membership_positive = ifelse(net_utility_positive >= 0, 1, 0)) %>% 
      # Calculate utilities for the negative world
      # Utility to join
      mutate(utility_join_negative = (-abs(climate_ideal_a - theta)) + # Climate
               (1 - share_trade_club_negative) * (1 - lambda_negative)*beta + # Non-club trade
               (share_trade_club_negative*beta) - # Club trade
               (theta - lambda_negative)) %>%  # Aversion
      # Utility to abstain
      mutate(utility_abstain_negative = (-climate_ideal_a) + # Climate
               (1-share_trade_club_negative)*beta + # Non-club trade
               (share_trade_club_negative*(1-lambda_negative)*beta) + # Club trade
               (0)) %>%  # Aversion
      # Net, membership
      mutate(net_utility_negative = utility_join_negative - utility_abstain_negative,
             membership_negative = ifelse(net_utility_negative >= 0, 1, 0)) %>% 
      # Define emissions common across states
      mutate(ghg_emissions = 100,
             carbon_price = 0.5*theta_draw) %>%
      mutate(delta_emissions_independent = ifelse(membership_independent==1,
                                                  ghg_emissions*(1-carbon_price),
                                                  ghg_emissions)) %>%
      mutate(delta_emissions_positive = ifelse(membership_positive==1,
                                               ghg_emissions*(1-carbon_price),
                                               ghg_emissions)) %>%
      mutate(delta_emissions_negative = ifelse(membership_negative==1,
                                               ghg_emissions*(1-carbon_price),
                                               ghg_emissions)) 
    
    # Save the number of countries in the club
    length_members_independent[which(theta_vector == theta_draw), 
                               which(seed == seeds_vector)] <- nrow(filter(theta_df, membership_independent == 1))
    length_members_positive[which(theta_vector == theta_draw), 
                            which(seed == seeds_vector)] <- nrow(filter(theta_df, membership_positive == 1))
    length_members_negative[which(theta_vector == theta_draw), 
                            which(seed == seeds_vector)] <- nrow(filter(theta_df, membership_negative == 1))
    
    # Save the emissions for values of theta
    length_emissions_independent[which(theta_vector == theta_draw), 
                                 which(seed == seeds_vector)] <- sum(pull(theta_df, delta_emissions_independent))
    length_emissions_positive[which(theta_vector == theta_draw), 
                              which(seed == seeds_vector)] <- sum(pull(theta_df, delta_emissions_positive))
    length_emissions_negative[which(theta_vector == theta_draw), 
                              which(seed == seeds_vector)] <- sum(pull(theta_df, delta_emissions_negative))
  }
}
end_time <- Sys.time()
end_time - start_time

#### Summarize outputs ####

## Summarize membership across iterations
members <- rbind(length_members_independent %>% 
                   as.data.frame() %>% 
                   mutate(theta = theta_vector) %>% 
                   relocate(theta, .before = "V1") %>% 
                   pivot_longer(names_to = "seed", values_to = "members", 2:ncol(.)) %>% 
                   group_by(theta) %>% 
                   summarise(mean_members = mean(members, na.rm = T)) %>% 
                   mutate(model = "Independent"),
                 length_members_negative %>% 
                   as.data.frame() %>% 
                   mutate(theta = theta_vector) %>% 
                   relocate(theta, .before = "V1") %>% 
                   pivot_longer(names_to = "seed", values_to = "members", 2:ncol(.)) %>% 
                   group_by(theta) %>% 
                   summarise(mean_members = mean(members, na.rm = T)) %>% 
                   mutate(model = "Negatively correlated"),
                 length_members_positive %>% 
                   as.data.frame() %>% 
                   mutate(theta = theta_vector) %>% 
                   relocate(theta, .before = "V1") %>% 
                   pivot_longer(names_to = "seed", values_to = "members", 2:ncol(.)) %>%  
                   group_by(theta) %>% 
                   summarise(mean_members = mean(members, na.rm = T)) %>% 
                   mutate(model = "Positively correlated")) 

## Summarize emissions across iterations
emissions <- rbind(length_emissions_independent %>% 
                     as.data.frame() %>% 
                     mutate(theta = theta_vector) %>% 
                     relocate(theta, .before = "V1") %>% 
                     pivot_longer(names_to = "seed", values_to = "emissions", 2:ncol(.)) %>% 
                     group_by(theta) %>% 
                     summarise(mean_emissions = mean(emissions, na.rm = T)) %>% 
                     mutate(model = "Independent"),
                   length_emissions_negative %>% 
                     as.data.frame() %>% 
                     mutate(theta = theta_vector) %>% 
                     relocate(theta, .before = "V1") %>% 
                     pivot_longer(names_to = "seed", values_to = "emissions", 2:ncol(.)) %>% 
                     group_by(theta) %>% 
                     summarise(mean_emissions = mean(emissions, na.rm = T)) %>% 
                     mutate(model = "Negatively correlated"),
                   length_emissions_positive %>% 
                     as.data.frame() %>% 
                     mutate(theta = theta_vector) %>% 
                     relocate(theta, .before = "V1") %>% 
                     pivot_longer(names_to = "seed", values_to = "emissions", 2:ncol(.)) %>% 
                     group_by(theta) %>% 
                     summarise(mean_emissions = mean(emissions, na.rm = T)) %>% 
                     mutate(model = "Positively correlated"))

## Bind members and emissions outputs and save to disk --- redundant, not run
# full_join(members, emissions, by = c("theta", "model")) %>%
#   write_csv(., "output/sims_broader_deeper_IxI.csv")

## Load in simulation output from disk
outs <- read_csv("output/sims_broader_deeper_IxI.csv")

## Separate out members and emissions summaries
members <- outs %>%
  select(theta, model, mean_members)
emissions <- outs %>%
  select(theta, model, mean_emissions)

## Plot membership
plot_members <- members %>% 
  ggplot(., aes(theta, mean_members, color = model)) +
  geom_line(size = 2) +
  geom_point(size = 4, color = "black") +
  geom_point(size = 3) +
  labs(x = expression(theta[j]),
       y = "Climate club members",
       color = "Model") +
  scale_x_continuous(limits = c(0,1), breaks=seq(0,1,0.1)) +
  scale_y_continuous(breaks=seq(0,150,25)) +
  scale_color_manual(values = c("#7DCEA0", "#229954", "#556B2F")) +
  theme_clubs +
  guides(col = guide_legend(nrow = 1))
plot_members
ggsave(plot_members, filename = "text/membership_across_theta.pdf", units = 'in', height = 6, width = 8)


## Plot emissions
plot_emissions <- emissions %>% 
  ggplot(., aes(theta, mean_emissions, color = model)) +
  coord_cartesian(ylim = c(12000,15000)) +
  geom_line(size = 2) +
  geom_point(size = 4, color = "black") +
  geom_point(size = 3) +
  labs(x = expression(theta[j]),
       y = "Aggregate GHG emissions",
       color = "Model") +
  scale_x_continuous(limits = c(0,1), breaks=seq(0,1,0.1)) +
  scale_color_manual(values = c("#7DCEA0", "#229954", "#556B2F")) +
  theme_clubs +
  guides(col = guide_legend(nrow = 1))
plot_emissions
ggsave(plot_emissions, filename = "text/emissions_across_theta.pdf", units = 'in', height = 6, width = 8)



